import time
import numpy as np
from PIL import Image

from encoder.clip import clip_encoder
from encoder.color import encode_color
from index_loader import clip_index_loader
from index_loader import color_index_loader
from index_search import search, mixture_of_experts_search

# load index
print("loading index")
clip_indices = clip_index_loader.load_index()
clip_invert_indices = clip_index_loader.load_invert_index()

color_indices = color_index_loader.load_index()
color_invert_indices = color_index_loader.load_invert_index()

experts = [
    (clip_indices, clip_invert_indices),
    (color_indices, color_invert_indices)
]

def search_by_text(text:str, topK=5):
    start_time = time.time()
    features = clip_encoder.encode_text(text)

    end_time = time.time()
    print(f"encode text time: {end_time - start_time}s")

    ret = search(features, clip_indices, clip_invert_indices, topK=topK)
    
    return ret

def search_by_img(img:Image, weights=[0.6, 0.4], topK=5):
    start_time = time.time()

    clip_features = clip_encoder.encode_img(img)
    color_features = encode_color(img)

    end_time = time.time()
    print(f"encode image time: {end_time - start_time}s")

    return mixture_of_experts_search([clip_features, color_features], experts, weights=weights, topK=topK)

def test_text():
    while True:
        query = input("输入查询：")
        if query == "quit":
            break
        matched_imgs = search_by_text(query, topK=5)
        for path, dist in matched_imgs:
            print("dist: ", dist)
            img = Image.open(path)
            img.show()

def test_img():
    query_img = Image.open("E:/ir/clip-base-image-retrieval-system/dog.jpeg")
    matched_imgs = search_by_img(query_img)
    for path, dist in matched_imgs:
        print("dist: ", dist)
        img = Image.open(path)
        img.show()

if __name__ == '__main__':
    #test_text()
    test_img()